/////////////////////////////////////////////////////////////////////////////
// Original code from libhdfs3. Copyright (c) 2013 - 2014, Pivotal Inc.
// All rights reserved. Author: Zhanwei Wang
/////////////////////////////////////////////////////////////////////////////
//  Modifications by Kumo Inc.
// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <ktest/ktest.h>
#include <kmock/kmock.h>

#include <kmhdfs/client/file_system.h>
#include <kmhdfs/client/file_system_impl.h>
#include <kmhdfs/client/file_system_inter.h>
#include <kmhdfs/client/output_stream_impl.h>
#include <kmhdfs/client/packet.h>
#include <kmhdfs/client/pipeline.h>
#include <tests/mock/mock_file_system_inter.h>
#include <tests/mock/mock_lease_renewer.h>
#include <tests/mock/mock_pipeline.h>
#include <tests/mock/namenode_stub.h>
#include <tests/mock/test_util.h>
#include <kmhdfs/common/xml_config.h>

#include <string>

using namespace Hdfs;
using namespace Internal;
using namespace Hdfs::Mock;
using namespace testing;
using ::testing::AtLeast;

#define BASE_DIR "test/"

class TestOutputStream: public ::testing::Test {
public:
    TestOutputStream() {
        renewer = MakeMockLeaseRenewer();
    }

    ~TestOutputStream() {
        ResetMockLeaseRenewer(renewer);
    }

protected:
    std::shared_ptr<LeaseRenewer> renewer;
    MockFileSystemInter fs;
    OutputStreamImpl ous;
};

class MockPipelineStub: public PipelineStub {
public:
    MOCK_METHOD0(getPipeline, std::shared_ptr<MockPipeline> ());
};

class MockNamenodeStub: public NamenodeStub {
public:
    MOCK_METHOD0(getNamenode, MockNamenode * ());
};

static void wrappedOpen(OutputStreamImpl & ous, std::shared_ptr<FileSystemInter> fs, const char * path,
                        int flag = Create, const Permission permission = Permission(0644),
                        bool createParent = false, int replication = 0, int64_t blockSize = 0) {
    std::pair<std::shared_ptr<LocatedBlock>, std::shared_ptr<Hdfs::FileStatus>> pair;
    if (flag & Append) {
        pair = fs->append(path, flag);
    } else {
        FileStatus status = fs->create(path, permission, flag, createParent, replication, blockSize);
        std::pair<std::shared_ptr<LocatedBlock>, std::shared_ptr<FileStatus>> retval;
        pair.first = std::shared_ptr<LocatedBlock>(nullptr);
        pair.second = std::shared_ptr<FileStatus>(new FileStatus(status));
    }
    std::shared_ptr<Hdfs::FileStatus> status = pair.second;
    ous.open(fs, path, pair, flag, 0777, false, replication, blockSize, 0);
}

static void LeaseRenew(int flag) {
    Config conf;
    FileStatus fileinfo;
    fileinfo.setBlocksize(2048);
    fileinfo.setLength(1024);
    std::shared_ptr<LocatedBlock> lastBlock(new LocatedBlock);
    lastBlock->setNumBytes(0);
    std::pair<std::shared_ptr<LocatedBlock>, std::shared_ptr<FileStatus> > lastBlockWithStatus;
    lastBlockWithStatus.first = lastBlock;
    lastBlockWithStatus.second = std::shared_ptr<FileStatus>(new FileStatus(fileinfo));
    MockNamenodeStub stub;
    SessionConfig sconf(conf);
    std::shared_ptr<MockFileSystemInter> myfs(new MockFileSystemInter());
    EXPECT_CALL(*myfs, getConf()).Times(1).WillOnce(ReturnRef(sconf));
    //EXPECT_CALL(stub, getNamenode()).Times(1).WillOnce(Return(nn));
    OutputStreamImpl leaseous;

    if (flag & Append) {
        EXPECT_CALL(*myfs, append(_,_)).Times(1);
    } else {
        EXPECT_CALL(*myfs, create(_, _, _, _, _, _)).Times(1);
    }

    EXPECT_CALL(*myfs, complete(_, _, _)).Times(1).WillOnce(Return(true));
    EXPECT_CALL(GetMockLeaseRenewer(), StartRenew(_)).Times(1);
    EXPECT_CALL(GetMockLeaseRenewer(), StopRenew(_)).Times(1);
    EXPECT_CALL(*myfs, getStandardPath(_)).Times(1);

    EXPECT_NO_THROW(DebugException(wrappedOpen(leaseous, myfs, BASE_DIR"testrenewlease", flag, 0644, true, 0, 2048)));
    EXPECT_NO_THROW(leaseous.close());
};

static void heartBeatSender(int flag) {
    OutputStreamImpl ous;
    std::shared_ptr<MockPipeline> pipeline(new MockPipeline());
    MockPipelineStub stub;
    ous.stub = &stub;
    MockFileSystemInter * fs = new MockFileSystemInter;
    FileStatus fileinfo;
    fileinfo.setBlocksize(2048);
    fileinfo.setLength(1024);
    Config conf;
    const SessionConfig sessionConf(conf);
    std::shared_ptr<LocatedBlock> lastBlock(new LocatedBlock);
    lastBlock->setNumBytes(0);
    std::pair<std::shared_ptr<LocatedBlock>, std::shared_ptr<FileStatus> > lastBlockWithStatus;
    lastBlockWithStatus.first = lastBlock;
    lastBlockWithStatus.second = std::shared_ptr<FileStatus>(new FileStatus(fileinfo));
    EXPECT_CALL(*fs, getStandardPath(_)).Times(1).WillOnce(Return("/testheartBeat"));
    EXPECT_CALL(*fs, getConf()).Times(1).WillOnce(ReturnRef(sessionConf));

    if (flag & Append) {
        EXPECT_CALL(*fs, append(_,_)).Times(1).WillOnce(Return(lastBlockWithStatus));
    } else {
        EXPECT_CALL(*fs, create(_, _, _, _, _, _)).Times(1);
    }

    EXPECT_CALL(*fs, registerOpenedOutputStream()).Times(1);
    EXPECT_NO_THROW(wrappedOpen(ous, std::shared_ptr<FileSystemInter>(fs), "testheartBeat", flag, 0644, false, 3, 1024 * 1024));
    char buffer[20];
    Hdfs::FillBuffer(buffer, sizeof(buffer), 0);
    EXPECT_NO_THROW(ous.append(buffer, sizeof(buffer)));
    EXPECT_CALL(stub, getPipeline()).Times(1).WillOnce(Return(pipeline));
    EXPECT_CALL(*pipeline, send(_)).Times(3);
    EXPECT_CALL(*pipeline, flush()).Times(1);
    EXPECT_NO_THROW(ous.flush());
    sleep_for(seconds(21));
    EXPECT_CALL(*pipeline, close(_)).Times(1).WillOnce(Return(lastBlock));
    EXPECT_CALL(*fs, fsync(_)).Times(1);
    EXPECT_CALL(*fs, complete(_, _, _)).Times(1).WillOnce(Return(true));
    EXPECT_CALL(*fs, unregisterOpenedOutputStream()).Times(1);
    EXPECT_NO_THROW(ous.close());
}

static void heartBeatSenderThrow(int flag) {
    OutputStreamImpl ous;
    std::shared_ptr<MockPipeline> pipeline(new MockPipeline());
    MockPipelineStub stub;
    ous.stub = &stub;
    MockFileSystemInter * fs = new MockFileSystemInter;
    FileStatus fileinfo;
    fileinfo.setBlocksize(2048);
    fileinfo.setLength(1024);
    Config conf;
    const SessionConfig sessionConf(conf);
    std::shared_ptr<LocatedBlock> lastBlock(new LocatedBlock);
    HdfsIOException e("test", "test", 3, "test");
    lastBlock->setNumBytes(0);
    std::pair<std::shared_ptr<LocatedBlock>, std::shared_ptr<FileStatus> > lastBlockWithStatus;
    lastBlockWithStatus.first = lastBlock;
    lastBlockWithStatus.second = std::shared_ptr<FileStatus>(new FileStatus(fileinfo));
    EXPECT_CALL(*fs, getStandardPath(_)).Times(1).WillOnce(Return("/testheartBeat"));
    EXPECT_CALL(*fs, getConf()).Times(1).WillOnce(ReturnRef(sessionConf));

    if (flag & Append) {
        EXPECT_CALL(*fs, append(_,_)).Times(1).WillOnce(Return(lastBlockWithStatus));
    } else {
        EXPECT_CALL(*fs, create(_, _, _, _, _, _)).Times(1);
    }

    EXPECT_CALL(*fs, registerOpenedOutputStream()).Times(1);
    EXPECT_NO_THROW(wrappedOpen(ous, std::shared_ptr<FileSystemInter>(fs), "testheartBeat", flag, 0644, false, 3, 1024 * 1024));
    char buffer[20];
    Hdfs::FillBuffer(buffer, sizeof(buffer), 0);
    EXPECT_NO_THROW(ous.append(buffer, sizeof(buffer)));
    EXPECT_CALL(stub, getPipeline()).Times(1).WillOnce(Return(pipeline));
    EXPECT_CALL(*pipeline, send(_)).Times(2).WillOnce(Return()).WillOnce(Throw(e));
    EXPECT_CALL(*pipeline, flush()).Times(1);
    EXPECT_NO_THROW(ous.flush());
    sleep_for(seconds(11));
    EXPECT_CALL(*pipeline, close(_)).Times(1).WillOnce(Return(lastBlock));
    EXPECT_CALL(*fs, fsync(_)).Times(1);
    EXPECT_CALL(*fs, complete(_, _, _)).Times(1).WillOnce(Return(true));
    EXPECT_CALL(*fs, unregisterOpenedOutputStream()).Times(1);
    EXPECT_NO_THROW(ous.close());
}

TEST_F(TestOutputStream, LeaseRenewForAppend_Success) {
    LeaseRenew(Create | Append);
}

TEST_F(TestOutputStream, LeaseRenewForCreate_Success) {
    LeaseRenew(Create);
}

TEST_F(TestOutputStream, DISABLED_heartBeatSenderForAppend_Success) {
    heartBeatSender(Create | Append);
}

TEST_F(TestOutputStream, DISABLED_heartBeatSenderForCreate_Success) {
    heartBeatSender(Create);
}

TEST_F(TestOutputStream, DISABLED_heartBeatSenderForCreate_Throw) {
    heartBeatSenderThrow(Create);
}

TEST_F(TestOutputStream, DISABLED_heartBeatSenderForAppend_Throw) {
    heartBeatSenderThrow(Create | Append);
}

TEST_F(TestOutputStream, openForCreate_Success) {
    OutputStreamImpl ous;
    MockFileSystemInter * fs = new MockFileSystemInter;
    Config conf;
    const SessionConfig sessionConf(conf);
    EXPECT_CALL(*fs, getStandardPath(_)).Times(1).WillOnce(Return("/testopen"));
    EXPECT_CALL(*fs, getConf()).Times(1).WillOnce(ReturnRef(sessionConf));
    EXPECT_CALL(*fs, create(_, _, _, _, _, _)).Times(1);
    EXPECT_CALL(GetMockLeaseRenewer(), StartRenew(_)).Times(1);
    EXPECT_CALL(GetMockLeaseRenewer(), StopRenew(_)).Times(1);
    EXPECT_NO_THROW(wrappedOpen(ous, std::shared_ptr<FileSystemInter>(fs), "testopen", Create, 0644, false, 3, 1024 * 1024));
    EXPECT_CALL(*fs, complete(_, _, _)).Times(1).WillOnce(Return(true));
    EXPECT_NO_THROW(ous.close());
}

TEST_F(TestOutputStream, registerForCreate_Success) {
    OutputStreamImpl ous;
    MockFileSystemInter * fs = new MockFileSystemInter;
    Config conf;
    const SessionConfig sessionConf(conf);
    EXPECT_CALL(*fs, getStandardPath(_)).Times(1).WillOnce(Return("/testregiester"));
    EXPECT_CALL(*fs, getConf()).Times(1).WillOnce(ReturnRef(sessionConf));
    EXPECT_CALL(*fs, create(_, _, _, _, _, _)).Times(1);
    EXPECT_CALL(GetMockLeaseRenewer(), StartRenew(_)).Times(1);
    EXPECT_CALL(GetMockLeaseRenewer(), StopRenew(_)).Times(1);
    EXPECT_NO_THROW(wrappedOpen(ous, std::shared_ptr<FileSystemInter>(fs), "testregiester", Create, 0644, false, 3, 1024 * 1024));
    EXPECT_CALL(*fs, complete(_, _, _)).Times(1).WillOnce(Return(true));
    EXPECT_NO_THROW(ous.close());
}

TEST_F(TestOutputStream, registerForAppend_Success) {
    OutputStreamImpl ous;
    MockFileSystemInter * fs = new MockFileSystemInter;
    Config conf;
    const SessionConfig sessionConf(conf);
    FileStatus fileinfo;
    fileinfo.setBlocksize(2048);
    fileinfo.setLength(1024);
    std::shared_ptr<LocatedBlock> lastBlock(new LocatedBlock);
    lastBlock->setNumBytes(0);
    std::pair<std::shared_ptr<LocatedBlock>, std::shared_ptr<FileStatus> > lastBlockWithStatus;
    lastBlockWithStatus.first = lastBlock;
    lastBlockWithStatus.second = std::shared_ptr<FileStatus>(new FileStatus(fileinfo));
    EXPECT_CALL(*fs, getStandardPath(_)).Times(1).WillOnce(Return("/testregiester"));
    EXPECT_CALL(*fs, getConf()).Times(1).WillOnce(ReturnRef(sessionConf));
    EXPECT_CALL(*fs, append(_,_)).Times(1);
    EXPECT_CALL(GetMockLeaseRenewer(), StartRenew(_)).Times(1);
    EXPECT_CALL(GetMockLeaseRenewer(), StopRenew(_)).Times(1);
    EXPECT_NO_THROW(wrappedOpen(ous, std::shared_ptr<FileSystemInter>(fs), "testregiester", Append, 0644, false, 0, 0));
    EXPECT_CALL(*fs, complete(_, _, _)).Times(1).WillOnce(Return(true));
    EXPECT_NO_THROW(ous.close());
}

TEST_F(TestOutputStream, openForCreate_Fail) {
    OutputStreamImpl ous;
    MockFileSystemInter * fs = new MockFileSystemInter;
    Config conf;
    const SessionConfig sessionConf(conf);
    HdfsIOException e("test", "test", 2, "test");
    EXPECT_CALL(*fs, create(_, _, _, _, _, _)).Times(1).WillOnce(Throw(e));
    EXPECT_THROW(wrappedOpen(ous, std::shared_ptr<FileSystemInter>(fs), "testopen", Create, 0644, false, 3, 1024 * 1024), HdfsIOException);
}


TEST_F(TestOutputStream, openForAppend_Success) {
    OutputStreamImpl ous;
    MockFileSystemInter * fs = new MockFileSystemInter;
    Config conf;
    const SessionConfig sessionConf(conf);
    FileStatus fileinfo;
    fileinfo.setBlocksize(2048);
    fileinfo.setLength(1024);
    std::shared_ptr<LocatedBlock> lastBlock(new LocatedBlock);
    lastBlock->setNumBytes(0);
    std::pair<std::shared_ptr<LocatedBlock>, std::shared_ptr<FileStatus> > lastBlockWithStatus;
    lastBlockWithStatus.first = lastBlock;
    lastBlockWithStatus.second = std::shared_ptr<FileStatus>(new FileStatus(fileinfo));
    EXPECT_CALL(*fs, getStandardPath(_)).Times(1).WillOnce(Return("/testopen"));
    EXPECT_CALL(*fs, getConf()).Times(1).WillOnce(ReturnRef(sessionConf));
    EXPECT_CALL(*fs, append(_,_)).Times(1).WillOnce(Return(lastBlockWithStatus));
    EXPECT_CALL(GetMockLeaseRenewer(), StartRenew(_)).Times(1);
    EXPECT_CALL(GetMockLeaseRenewer(), StopRenew(_)).Times(1);
    EXPECT_NO_THROW(wrappedOpen(ous, std::shared_ptr<FileSystemInter>(fs), "testopen", Append, 0644, false, 0, 0));
    EXPECT_CALL(*fs, complete(_, _, _)).Times(1).WillOnce(Return(true));
    EXPECT_NO_THROW(ous.close());
}

TEST_F(TestOutputStream, openForAppend_Fail) {
    OutputStreamImpl ous;
    MockFileSystemInter * fs = new MockFileSystemInter;
    Config conf;
    const SessionConfig sessionConf(conf);
    FileStatus fileinfo;
    fileinfo.setBlocksize(2048);
    fileinfo.setLength(1024);
    EXPECT_CALL(*fs, append(_,_)).Times(1).WillOnce(Throw(FileNotFoundException("test", "test", 2, "test")));
    EXPECT_THROW(wrappedOpen(ous, std::shared_ptr<FileSystemInter>(fs), "testopen", Append, 0644, false, 0, 0), FileNotFoundException);
}

TEST_F(TestOutputStream, append_Success) {
    OutputStreamImpl ous;
    std::shared_ptr<MockPipeline> pipelineStub(new MockPipeline());
    MockPipelineStub stub;
    ous.stub = &stub;
    MockFileSystemInter * fs = new MockFileSystemInter;
    FileStatus fileinfo;
    fileinfo.setBlocksize(2048);
    fileinfo.setLength(1024);
    Config conf;
    const SessionConfig sessionConf(conf);
    std::shared_ptr<LocatedBlock> lastBlock(new LocatedBlock);
    lastBlock->setNumBytes(0);
    std::pair<std::shared_ptr<LocatedBlock>, std::shared_ptr<FileStatus> > lastBlockWithStatus;
    lastBlockWithStatus.first = lastBlock;
    lastBlockWithStatus.second = std::shared_ptr<FileStatus>(new FileStatus(fileinfo));
    EXPECT_CALL(*fs, getStandardPath(_)).Times(1).WillOnce(Return("/testopen"));
    EXPECT_CALL(*fs, getConf()).Times(1).WillOnce(ReturnRef(sessionConf));
    EXPECT_CALL(*fs, append(_,_)).Times(1).WillOnce(Return(lastBlockWithStatus));
    EXPECT_CALL(GetMockLeaseRenewer(), StartRenew(_)).Times(1);
    EXPECT_CALL(GetMockLeaseRenewer(), StopRenew(_)).Times(1);
    EXPECT_NO_THROW(wrappedOpen(ous, std::shared_ptr<FileSystemInter>(fs), "testopen", Create | Append, 0644, false, 3, 2048));
    char buffer[4096 + 523];
    Hdfs::FillBuffer(buffer, sizeof(buffer), 0);
    EXPECT_CALL(stub, getPipeline()).Times(3).WillOnce(Return(pipelineStub)).WillOnce(Return(pipelineStub)).WillOnce(Return(pipelineStub));
    EXPECT_CALL(*pipelineStub, send(_)).Times(4);
    EXPECT_CALL(*pipelineStub, close(_)).Times(2).WillOnce(Return(lastBlock)).WillOnce(Return(lastBlock));
    EXPECT_CALL(*fs, fsync(_)).Times(2);
    EXPECT_NO_THROW(ous.append(buffer, sizeof(buffer)));
    EXPECT_CALL(*pipelineStub, close(_)).Times(1).WillOnce(Return(lastBlock));
    EXPECT_CALL(*fs, fsync(_)).Times(1);
    EXPECT_CALL(*fs, complete(_, _, _)).Times(1).WillOnce(Return(true));
    EXPECT_NO_THROW(ous.close());
}

TEST_F(TestOutputStream, flush_Success) {
    OutputStreamImpl ous;
    std::shared_ptr<MockPipeline> pipelineStub(new MockPipeline());
    MockPipelineStub stub;
    ous.stub = &stub;
    MockFileSystemInter * fs = new MockFileSystemInter;
    FileStatus fileinfo;
    fileinfo.setBlocksize(2048);
    fileinfo.setLength(1024);
    Config conf;
    const SessionConfig sessionConf(conf);
    std::shared_ptr<LocatedBlock> lastBlock(new LocatedBlock);
    HdfsIOException e("test", "test", 3, "test");
    lastBlock->setNumBytes(0);
    std::pair<std::shared_ptr<LocatedBlock>, std::shared_ptr<FileStatus> > lastBlockWithStatus;
    lastBlockWithStatus.first = lastBlock;
    lastBlockWithStatus.second = std::shared_ptr<FileStatus>(new FileStatus(fileinfo));
    EXPECT_CALL(*fs, getStandardPath(_)).Times(1).WillOnce(Return("/testflush"));
    EXPECT_CALL(*fs, getConf()).Times(1).WillOnce(ReturnRef(sessionConf));
    EXPECT_CALL(*fs, append(_,_)).Times(1).WillOnce(Return(lastBlockWithStatus));
    EXPECT_CALL(GetMockLeaseRenewer(), StartRenew(_)).Times(1);
    EXPECT_CALL(GetMockLeaseRenewer(), StopRenew(_)).Times(1);
    EXPECT_NO_THROW(wrappedOpen(ous, std::shared_ptr<FileSystemInter>(fs), "testflush", Create | Append, 0644, false, 3, 1024 * 1024));
    char buffer[20];
    Hdfs::FillBuffer(buffer, sizeof(buffer), 0);
    EXPECT_NO_THROW(ous.append(buffer, sizeof(buffer)));
    EXPECT_CALL(stub, getPipeline()).Times(1).WillOnce(Return(pipelineStub));
    EXPECT_CALL(*pipelineStub, send(_)).Times(1);
    EXPECT_CALL(*pipelineStub, flush()).Times(1);
    EXPECT_NO_THROW(ous.flush());
    EXPECT_CALL(*pipelineStub, close(_)).Times(1).WillOnce(Return(lastBlock));
    EXPECT_CALL(*fs, fsync(_)).Times(1);
    EXPECT_CALL(*fs, complete(_, _, _)).Times(1).WillOnce(Return(true));
    EXPECT_NO_THROW(ous.close());
}

TEST_F(TestOutputStream, ValidateFirstBadLink) {
    EXPECT_NO_THROW(PipelineImpl::checkBadLinkFormat(""));
    EXPECT_NO_THROW(PipelineImpl::checkBadLinkFormat("8.8.8.8:1234"));
    EXPECT_NO_THROW(PipelineImpl::checkBadLinkFormat("2001:0db8:85a3:0000:0000:8a2e:0370:7334:50010"));
    EXPECT_THROW(PipelineImpl::checkBadLinkFormat("3"), HdfsException);
    EXPECT_THROW(PipelineImpl::checkBadLinkFormat("8.8.8.8"), HdfsException);
    EXPECT_THROW(PipelineImpl::checkBadLinkFormat("8.8.8.8:"), HdfsException);
    EXPECT_THROW(PipelineImpl::checkBadLinkFormat("8.8.8.888:50010"), HdfsException);
    EXPECT_THROW(PipelineImpl::checkBadLinkFormat("8.8.8.8:500101"), HdfsException);
    EXPECT_THROW(PipelineImpl::checkBadLinkFormat("8.8.8.8:50010a"), HdfsException);
}
